package com.atguigu1.core.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

/**
 * @description: udaf
 * @time: 2021/3/25 11:28
 * @author: baojinlong
 **/
object SparkSqlDemo03 {
  def main(args: Array[String]): Unit = {
    // 创建SparkSql的运行环境
    val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("sparkSql")
    val spark: SparkSession = SparkSession.builder.config(sparkConf).getOrCreate

    val dataFrame: DataFrame = spark.read.json("datas/user.json")
    dataFrame.createOrReplaceTempView("user")
    // spark.sql("select age,username from user").show
    // 注册函数
    // spark.udf.register("prefixName", (name: String) => "Name:" + name)
    spark.udf.register("ageAvg", new MyAvgAgeUdaf)
    spark.sql("select ageAvg(age) from user").show
    // 关闭环境
    spark.close
  }

  /**
   * 计算年龄平均值
   */
  class MyAvgAgeUdaf extends UserDefinedAggregateFunction {
    // 返回的结构类型
    override def inputSchema: StructType = {
      StructType(Array(StructField("age", LongType)))
    }

    /**
     * 缓冲区数据的结构
     *
     * @return
     */
    override def bufferSchema: StructType = {
      StructType(
        Array(
          StructField("total", LongType),
          StructField("count", LongType)
        )
      )
    }

    // 函数计算结果的数据类型out
    override def dataType: DataType = LongType

    // 函数的稳定性
    override def deterministic: Boolean = true

    // 缓冲区初始化
    override def initialize(buffer: MutableAggregationBuffer): Unit = {
      // buffer(0,0L)
      // buffer(1,0L)
      buffer.update(0, 0L)
      buffer.update(1, 0L)
    }


    // 根据输入的值更新缓冲区数据
    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
      // 更新年龄总和,更新数据个数
      buffer.update(0, buffer.getLong(0) + input.getLong(0))
      buffer.update(1, buffer.getLong(1) + 1)
    }

    // 缓冲区数据合并
    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
      buffer1.update(0, buffer1.getLong(0) + buffer2.getLong(0))
      buffer1.update(1, buffer1.getLong(1) + buffer2.getLong(1))
    }

    // 计算平均值
    override def evaluate(buffer: Row): Any = {
      buffer.getLong(0) / buffer.getLong(1)
    }
  }

}
